from rdkit import Chem

def load_smiles(file_path):
    """
    从文件中加载 SMILES 分子列表
    :param file_path: 文本文件路径，包含生成的分子
    :return: SMILES 列表
    """
    smiles_list = []
    with open(file_path, "r") as f:
        for line in f:
            if ":" in line:  # 假设 SMILES 分子格式为 "Molecule X: SMILES"
                smiles = line.split(":")[1].strip()
                smiles_list.append(smiles)
    return smiles_list


def is_valid_smiles(smiles):
    """
    检查 SMILES 是否有效
    :param smiles: SMILES 字符串
    :return: 如果有效返回 True，否则返回 False
    """
    try:
        mol = Chem.MolFromSmiles(smiles)  # 转换为 RDKit 分子对象
        return mol is not None  # 如果分子对象有效，则返回 True
    except:
        return False


def calculate_validity(smiles_list):
    """
    计算 SMILES 分子的有效性
    :param smiles_list: SMILES 字符串列表
    :return: 有效性百分比
    """
    total_count = len(smiles_list)
    valid_count = sum(is_valid_smiles(smiles) for smiles in smiles_list)

    # 打印详细信息
    print(f"总分子数: {total_count}")
    print(f"有效分子数: {valid_count}")
    print(f"无效分子数: {total_count - valid_count}")

    # 计算有效性
    validity = (valid_count / total_count) * 100
    return validity


def save_valid_molecules(smiles_list, output_file):
    """
    保存有效的分子到文件
    :param smiles_list: SMILES 字符串列表
    :param output_file: 输出文件路径
    """
    with open(output_file, "w") as f:
        for smiles in smiles_list:
            if is_valid_smiles(smiles):
                f.write(smiles + "\n")
    print(f"有效分子已保存到 {output_file}")


if __name__ == "__main__":
    # 指定生成分子文件路径
    input_file = "./generated_molecules_full.txt"  # 假设生成分子保存为 generated_molecules.txt
    output_file = "./generated_molecules_effective.txt"  # 保存有效分子的文件

    # 加载生成的 SMILES
    print("加载生成的分子...")
    smiles_list = load_smiles(input_file)

    # 仅选择前 100 个分子
    smiles_list = smiles_list[:100]

    # 计算有效性
    print("计算分子的有效性...")
    validity = calculate_validity(smiles_list)

    # 打印结果
    print(f"分子有效性: {validity:.2f}%")

    # 保存有效分子到文件
    save_valid_molecules(smiles_list, output_file)
